"""
evacuation_maze_fixed.py

Advanced 3D evacuation simulation with:
- Structured maze-like floor layouts with continuous walls
- Adjustable path widths (1-3 cells)
- Multiple staircases per floor
- Guaranteed path connectivity
- Realistic crowd dynamics
"""

import numpy as np
import heapq
import math
import random
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
from scipy.ndimage import binary_dilation, binary_erosion
from collections import deque

# -----------------------------
# 1. 参数设置
# -----------------------------
CELL_SIZE         = 1.0    # 元胞边长 (m)
FLOORS            = 10     # 楼层数
HEIGHT, WIDTH     = 30, 30 # 每层网格尺寸
AGENTS_PER_FLOOR  = 40     # 每层人数
SPEED_MEAN        = 1.0    # m/s，地面初速度均值
SPEED_STD         = 0.1    # m/s，速度标准差
STAIR_SPEED_FACTOR= 0.7    # 楼梯移动速度相对地面速度系数
FLOOR_HEIGHT      = 3.0    # 每层楼高 (m)
MAX_STEPS         = 1000   # 最大模拟步数
STAIR_COUNT       = 3      # 每层楼梯数量
OBSTACLE_DENSITY  = 0.25   # 迷宫墙体密度（占总面积的25%）
ANIMATION_SPEED   = 500    # 动画间隔时间（毫秒）

# -----------------------------
# 2. 迷宫生成算法（递归回溯改进版）
# -----------------------------
def generate_maze(height, width):
    cell_size = 2  # 每个单元格的大小
    maze_size = (height // cell_size, width // cell_size)
    maze = np.ones((height, width), dtype=int)
    visited = np.zeros((maze_size[0], maze_size[1]), dtype=bool)
    
    def carve_maze(y, x):
        dirs = [(-1,0),(1,0),(0,-1),(0,1)]
        random.shuffle(dirs)
        
        for dy, dx in dirs:
            ny_cell = y + dy
            nx_cell = x + dx
            if 0 <= ny_cell < maze_size[0] and 0 <= nx_cell < maze_size[1]:
                if not visited[ny_cell, nx_cell]:
                    visited[ny_cell][nx_cell] = True
                    
                    # 计算墙壁拆除的坐标范围
                    current_y = y * cell_size
                    current_x = x * cell_size
                    neighbor_y = ny_cell * cell_size
                    neighbor_x = nx_cell * cell_size
                    
                    # 垂直方向拆除
                    if dy != 0:
                        sy = current_y + dy * (cell_size // 2)
                        ey = neighbor_y - dy * (cell_size // 2)
                        for i in range(cell_size // 2):
                            maze[sy + i][current_x] = 0
                            maze[ey - i][neighbor_x] = 0
                    # 水平方向拆除
                    else:
                        sx = current_x + dx * (cell_size // 2)
                        ex = neighbor_x - dx * (cell_size // 2)
                        for i in range(cell_size // 2):
                            maze[current_y][sx + i] = 0
                            maze[neighbor_y][ex - i] = 0
                    
                    carve_maze(ny_cell, nx_cell)
    
    visited[0, 0] = True
    carve_maze(0, 0)
    
    # 添加随机宽度通道
    for _ in range(random.randint(1, 3)):
        direction = random.choice(['horizontal', 'vertical', 'both'])
        if direction == 'horizontal':
            structure = np.array([[0,1,0],[0,1,0],[0,1,0]])
        elif direction == 'vertical':
            structure = np.array([[0,0,0],[1,1,1],[0,0,0]])
        else:
            structure = np.ones((3,3))
        
        maze = binary_dilation(maze == 0, structure=structure, iterations=1).astype(int)
    
    # 添加出口和入口
    maze[1, 1] = 0
    maze[-2, -2] = 0
    
    return maze

# -----------------------------
# 3. 楼层布局生成（改进版）
# -----------------------------
def build_environment():
    env = np.ones((FLOORS, HEIGHT, WIDTH), dtype=int)
    stair_positions_per_floor = []
    
    for z in range(FLOORS):
        maze = generate_maze(HEIGHT, WIDTH)
        
        # 添加楼梯
        stair_positions = []
        for _ in range(STAIR_COUNT):
            while True:
                y = random.randint(2, HEIGHT-3)
                x = random.randint(2, WIDTH-3)
                if maze[y][x] == 0:
                    stair_positions.append((y, x))
                    break
        
        # 楼梯区域处理
        for (sy, sx) in stair_positions:
            maze[sy-1:sy+2, sx-1:sx+2] = 0
            maze[sy, sx] = 2  # 楼梯标记
        
        # 确保连通性
        for (sy, sx) in stair_positions:
            bfs(maze, sy, sx, stair_positions)
        
        env[z] = maze
        stair_positions_per_floor.append(stair_positions)
    
    exit_positions = []
    for (sy, sx) in stair_positions_per_floor[0]:
        exit_positions.append((0, sy, sx))
    
    return env, stair_positions_per_floor, exit_positions

# -----------------------------
# 4. BFS 连通性检查（改进版）
# -----------------------------
def bfs(grid, sy, sx, stair_positions):
    visited = np.zeros_like(grid, dtype=bool)
    q = deque()
    q.append((sy, sx))
    visited[sy][sx] = True
    
    dirs = [(-1,0),(1,0),(0,-1),(0,1)]
    
    while q:
        y, x = q.popleft()
        for dy, dx in dirs:
            ny, nx = y + dy, x + dx
            if 0 <= ny < HEIGHT and 0 <= nx < WIDTH:
                if not visited[ny][nx] and grid[ny][nx] != 1:
                    visited[ny][nx] = True
                    q.append((ny, nx))
    
    # 检查未连通区域
    unvisited = np.where((grid == 0) & ~visited)
    for y, x in zip(*unvisited):
        shortest = None
        for (s_y, s_x) in stair_positions:
            dist = math.hypot(y - s_y, x - s_x)
            if shortest is None or dist < shortest[0]:
                shortest = (dist, (s_y, s_x))
        target = shortest[1]
        path = a_star(grid, (y, x), target)
        for step in path:
            grid[step] = 0

# -----------------------------
# 5. A* 寻路算法（保持原逻辑）
# -----------------------------
def a_star(grid, start, goal):
    open_set = []
    heapq.heappush(open_set, (0, start))
    came_from = {}
    g_score = {start:0}
    f_score = {start:heuristic(start, goal)}
    
    dirs = [(-1,0),(1,0),(0,-1),(0,1)]
    
    while open_set:
        current = heapq.heappop(open_set)[1]
        if current == goal:
            return reconstruct_path(came_from, current)
        for dy, dx in dirs:
            neighbor = (current[0]+dy, current[1]+dx)
            if (0 <= neighbor[0] < HEIGHT and 
                0 <= neighbor[1] < WIDTH and 
                grid[neighbor[0]][neighbor[1]] != 1):
                tentative_g = g_score[current] + 1
                if (neighbor not in g_score or 
                    tentative_g < g_score.get(neighbor, np.inf)):
                        came_from[neighbor] = current
                        g_score[neighbor] = tentative_g
                        f_score[neighbor] = tentative_g + heuristic(neighbor, goal)
                        heapq.heappush(open_set, (f_score[neighbor], neighbor))
    return []

def heuristic(a, b):
    return math.hypot(a[0]-b[0], a[1]-b[1])

def reconstruct_path(came_from, current):
    total_path = [current]
    while current in came_from:
        current = came_from[current]
        total_path.append(current)
    return total_path[::-1]

# -----------------------------
# 6. 静态场计算（保持原逻辑）
# -----------------------------
def compute_static_field(env, exits):
    f, h, w = env.shape
    static = np.full((f, h, w), np.inf)
    visited = np.zeros((f, h, w), dtype=bool)
    pq = []
    
    for pos in exits:
        static[pos] = 0.0
        heapq.heappush(pq, (0.0, pos))
    
    dirs = [(dz, dy, dx) 
            for dz in (-1,0,1) 
            for dy in (-1,0,1) 
            for dx in (-1,0,1) 
            if not (dz==dy==dx==0)]
    
    while pq:
        dist_curr, (z, y, x) = heapq.heappop(pq)
        if visited[z, y, x]:
            continue
        visited[z, y, x] = True
        
        for dz, dy, dx in dirs:
            nz, ny, nx = z+dz, y+dy, x+dx
            if (0 <= nz < f and 0 <= ny < h and 0 <= nx < w and 
                env[nz, ny, nx] != 1):
                step = (CELL_SIZE * math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) 
                        if dz != 0 else CELL_SIZE)
                nd = dist_curr + step
                if nd < static[nz, ny, nx]:
                    static[nz, ny, nx] = nd
                    heapq.heappush(pq, (nd, (nz, ny, nx)))
    return static

# -----------------------------
# 7. Agent 类（保持原逻辑）
# -----------------------------
class Agent:
    def __init__(self, z, y, x):
        self.pos = (z, y, x)
        self.V0 = random.gauss(SPEED_MEAN, SPEED_STD)
        self.d_sum = 0.0  # 累计移动距离
        self.t_sum = 0.0  # 累计时间
        self.count = 0.0
    
    def speed(self, dz):
        base = self.V0 * (1 - 0.01 * (self.d_sum**0.71) / 
                   (1 - 0.01 * (self.t_sum**0.51) * 
                    math.exp(0.01*self.t_sum)))
        return base * (STAIR_SPEED_FACTOR if dz != 0 else 1.0)

# -----------------------------
# 8. 初始化 Agents（保持原逻辑）
# -----------------------------
def init_agents(env):
    agents = []
    for z in range(FLOORS):
        empties = list(zip(*np.where(env[z] == 0)))
        random.shuffle(empties)
        for y, x in empties[:AGENTS_PER_FLOOR]:
            agents.append(Agent(z, y, x))
    return agents

# -----------------------------
# 9. 模拟循环（保持原逻辑）
# -----------------------------
def simulate(env, static, exits):
    agents = init_agents(env)
    frames = []
    stats = []
    
    dirs = [(dz, dy, dx) 
            for dz in (-1,0,1) 
            for dy in (-1,0,1) 
            for dx in (-1,0,1) 
            if not (dz==dy==dx==0)]
    
    for step in range(MAX_STEPS):
        occ = set(a.pos for a in agents)
        positions = []
        
        for a in agents[:]:
            a.t_sum += 1
            a.count += 1
            z, y, x = a.pos
            tau = CELL_SIZE / a.speed(0)
            
            if a.count >= tau:
                candidates = []
                for dz, dy, dx in dirs:
                    nz, ny, nx = z+dz, y+dy, x+dx
                    if (0 <= nz < FLOORS and 0 <= ny < HEIGHT and 0 <= nx < WIDTH and
                        env[nz, ny, nx] != 1 and 
                        (nz, ny, nx) not in occ and 
                        static[nz, ny, nx] <= static[z, y, x]):
                        candidates.append((static[nz, ny, nx], (dz, dy, dx)))
                
                if candidates:
                    candidates.sort(key=lambda x: x[0])
                    best_dist = candidates[0][0]
                    best_moves = [mv for d, mv in candidates if d == best_dist]
                    dz, dy, dx = random.choice(best_moves)
                    
                    new_pos = (z+dz, y+dy, x+dx)
                    occ.remove(a.pos)
                    occ.add(new_pos)
                    a.pos = new_pos
                    
                    dist = (CELL_SIZE * math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) 
                            if dz != 0 else CELL_SIZE)
                    a.d_sum += dist
                    a.count = 0
                    
                    if new_pos in exits:
                        occ.remove(new_pos)
                        agents.remove(a)
                else:
                    a.count = 0  # 无法移动则重置计数器
            
            positions.append(a.pos)
        
        total = len(agents)
        per_floor = [sum(1 for p in positions if p[0] == f) for f in range(FLOORS)]
        stats.append((total, per_floor))
        frames.append(positions)
        
        if not agents:
            break
    
    return frames, stats

# -----------------------------
# 10. 可视化（调整动画速度）
# -----------------------------
def visualize(env, frames, stats, stair_positions, exits):
    cmap = ListedColormap(['white', 'dimgray', 'saddlebrown'])
    fig = plt.figure(figsize=(20, 10))
    gs = fig.add_gridspec(2, 6)
    
    # 3D视图
    ax3d = fig.add_subplot(gs[:, 0], projection='3d')
    ax3d.set_box_aspect([WIDTH, HEIGHT, FLOORS])
    ax3d.set_xlabel('X')
    ax3d.set_ylabel('Y')
    ax3d.set_zlabel('Floor')
    
    # 绘制楼梯结构
    stairs = (env == 2)
    ax3d.voxels(stairs, facecolors='saddlebrown', 
                edgecolor='k', alpha=0.7)
    
    scat3d = ax3d.scatter([], [], [], s=20, c='blue')
    text_total = ax3d.text2D(0.02, 0.95, '', transform=ax3d.transAxes)
    
    # 2D楼层视图
    axs = []
    scats = []
    for f in range(FLOORS):
        row = 0 if f < 5 else 1
        col = (f % 5) + 1
        ax = fig.add_subplot(gs[row, col])
        ax.imshow(env[f], cmap=cmap, origin='lower')
        ax.set_title(f'Floor {f+1}')
        ax.axis('off')
        # 标记楼梯位置
        for sy, sx in stair_positions[f]:
            ax.scatter(sx, sy, c='red', s=50, marker='s')
        axs.append(ax)
        scats.append(ax.scatter([], [], s=10, c='red'))
    
    def update(frame):
        pts = frames[frame]
        total, per_floor = stats[frame]
        
        # 更新3D视图
        xs, ys, zs = [], [], []
        for (z, y, x) in pts:
            xs.append(x)
            ys.append(y)
            zs.append(z)
        scat3d._offsets3d = (xs, ys, zs)
        text_total.set_text(f'总剩余: {total}')
        
        # 更新各楼层视图
        for f in range(FLOORS):
            floor_agents = [(y, x) for (z, y, x) in pts if z == f]
            if floor_agents:
                ys_f, xs_f = zip(*floor_agents)
                scats[f].set_offsets(np.column_stack([xs_f, ys_f]))
            else:
                scats[f].set_offsets(np.empty((0, 2)))  # 空二维数组
        
        return [scat3d] + scats + [text_total]
    
    anim = animation.FuncAnimation(
        fig, update, frames=len(frames), 
        interval=ANIMATION_SPEED, blit=False
    )
    plt.show()

# -----------------------------
# 11. 主程序（保持原逻辑）
# -----------------------------
def main():
    env, stair_positions, exits = build_environment()
    static = compute_static_field(env, exits)
    frames, stats = simulate(env, static, exits)
    visualize(env, frames, stats, stair_positions, exits)

if __name__ == '__main__':
    main()